"""
=========================================================
HHL-style 相位估计（Quantum Phase Estimation, QPE）
=========================================================

本模块实现了 HHL 算法中用于提取特征相位（即特征值 λ 对应相位 φ=λt/2π）
的量子相位估计部分。该版本基于：
- cqlib 的 Algorithm_Circuit（量子电路对象）
- qft()（量子傅里叶变换）
- Pauli 演化 Oracle（实现 exp(-i * t * H) 的受控算子）

------------------------------------------
QPE 原理概述：
------------------------------------------
目标：估计满足
    U |u⟩ = e^{2πiφ} |u⟩
的酉算符 U 的特征相位 φ。

核心流程：
1. 准备 m 位“相位寄存器” + n 位“工作寄存器”（存储 |u⟩）。
2. 在相位寄存器上施加 H^{⊗m}；
3. 依次执行受控 U^{2^k}（第 k 位控制对应 2^k 次幂）；
4. 对相位寄存器施加逆 QFT（IQFT）；
5. 测量相位寄存器，得到 φ ≈ λ t / (2π)。

------------------------------------------
在 HHL 中的应用：
------------------------------------------
- U = exp(-i * H * t)
- 相位寄存器测得的 φ 代表 λt / 2π
- 可进一步根据 φ 构造 |1/λ⟩ 权重态，实现线性方程求解。

------------------------------------------
约定与实现说明：
------------------------------------------
- 相位寄存器 qubits 顺序为 LSB→MSB；
  第 k 位对应 U^{2^k} 的受控作用；
- IQFT 默认 `do_swaps=False`（不交换顺序）；
- HHL-style 相位估计通常不做最终比特交换。

"""

from __future__ import annotations
from typing import Callable, Iterable, Optional, Sequence, Tuple, Any, List
import numpy as np
from numpy.linalg import eigvalsh
from math import pi

# cqlib 电路类
from circuits.algorithm_circuit import Algorithm_Circuit as Circuit
# 你在 algorithms/qft.py 中实现的 QFT 接口
from algorithm.quantum_fourier_transform import qft


# ============================================================
# 类型别名：EvolutionOracle
# ============================================================
# 表示一个受控哈密顿演化操作：
# evolution_oracle(qc, ctrl, work, tau)
# 实现：U(tau) = exp(-i * tau * H)
EvolutionOracle = Callable[[Circuit, int, Iterable[int], float], None]


# ============================================================
# 相位估计：编码阶段（正向）
# ============================================================
def hhl_pe_encode(qc: Circuit,
                  phase_qubits: Iterable[int],
                  work_qubits: Iterable[int],
                  evolution_oracle: EvolutionOracle,
                  t: float) -> Circuit:
    """
    HHL 相位估计的“编码阶段”：
        H^m → [Controlled U(2^k t)]_k → IQFT(do_swaps=False)

    参数：
        qc                : Circuit 电路对象
        phase_qubits      : 相位寄存器量子比特索引（LSB→MSB）
        work_qubits       : 工作寄存器（被作用的系统态）
        evolution_oracle  : 受控时间演化算符 U(τ)=exp(-iτH)
        t                 : 基础时间步长（总演化时间的基本单位）

    实现步骤：
        1. 对相位寄存器施加 H（创建均匀叠加）；
        2. 逐位执行 Controlled-U^{2^k}；
        3. 对相位寄存器施加逆 QFT（IQFT）。

    说明：
        - 采用小端（LSB）为第一控制位；
        - IQFT 不执行 swap，保持低位在左、高位在右。
    """
    ph = list(phase_qubits)
    wrk = list(work_qubits)

    # 1) Hadamard 铺平相位寄存器
    for q in ph:
        qc.h(q)

    m = len(ph)
    # 2) 对第 k 位施加受控 U^{2^k}（幂次按 LSB→MSB 递增）
    for k, ctrl in enumerate(ph):
        power = 1 << k     # 2^(m-1-k)
        evolution_oracle(qc, ctrl, wrk, power * t)
    qc.barrier(0)
    # 3) IQFT on phase register
    qft(qc, ph, inverse=True, do_swaps=False)
    return qc


# ============================================================
# 相位估计：解码阶段（回滚）
# ============================================================
def hhl_pe_uncompute(qc: Circuit,
                     phase_qubits: Iterable[int],
                     work_qubits: Iterable[int],
                     evolution_oracle: EvolutionOracle,
                     t: float) -> Circuit:
    """
    HHL 相位估计的“解码阶段”（或称 uncompute）：
        QFT(do_swaps=False) → [Controlled U(2^k t)]† → H^m

    作用：
        撤销 hhl_pe_encode 所施加的所有操作，回到原始寄存器状态。

    参数：
        qc, phase_qubits, work_qubits, evolution_oracle, t
        （与 hhl_pe_encode 相同）

    步骤：
        1. 对相位寄存器施加正向 QFT；
        2. 按逆序依次执行受控 U^{2^k}†（取负角度）；
        3. 最后撤销所有 H。
    """
    ph = list(phase_qubits)
    wrk = list(work_qubits)

    # 1) 正向 QFT（不做 swap）
    qft(qc, ph, inverse=False, do_swaps=False)

    m = len(ph)
    # 2) 按逆序撤销所有受控演化（与 encode 完全互逆）
    for k, ctrl in reversed(list(enumerate(ph))):
        power = 1 << k
        evolution_oracle(qc, ctrl, wrk, - power * t)

    # 3) 撤销 H
    for q in ph:
        qc.h(q)
    return qc


# ============================================================
# 时间步长 t 的自动选择（根据谱界 λ_max）
# ============================================================
def pick_t_from_lambda_max(lambda_max: float,
                           target_phase: float = 0.35,
                           safety: float = 0.95) -> float:
    """
    根据最大特征值 λ_max 自动选择合适的时间步长 t。

    思路：
        使得 λ_max * t / (2π) ≈ target_phase（一般取 1/4）

    参数：
        lambda_max  : 最大特征值（谱半径上界）
        target_phase: 目标相位（默认 0.25）
        safety      : 安全系数（<1，稍微缩小 t 以防溢出）

    返回：
        推荐的时间步长 t。

    常用设置：
        - target_phase = 0.25 (HHL 常用)
        - safety = 0.95
    """
    if lambda_max <= 0:
        raise ValueError("lambda_max must be > 0 (SPD assumed in HHL).")
    return safety * (2.0 * pi * target_phase) / float(lambda_max)


def rough_lambda_max_via_eigs(H: np.ndarray) -> float:
    """
    粗略估计 H 的最大特征值 λ_max。

    小矩阵：直接调用 numpy.linalg.eigvalsh；
    大矩阵：建议使用 Gershgorin 圆盘定理或幂法近似。

    返回：
        float 最大特征值的绝对值。
    """
    w = eigvalsh(H)
    return float(np.max(np.abs(w)))


# ============================================================
# 工具函数：态矢量格式标准化与相位寄存器概率提取
# ============================================================
def _state_to_array(psi: Any, n: int) -> np.ndarray:
    """
    将各种格式（dict / ndarray / list）统一转为复数组。
    兼容：
        - dict: {bitstring: amplitude}
        - dict: {int_index: amplitude}
        - ndarray: 已是复数组
    """
    if isinstance(psi, dict):
        keys = list(psi.keys())
        if not keys:
            return np.zeros(2**n, dtype=complex)
        k0 = next(iter(keys))
        # 索引为整数的情况
        if isinstance(k0, (int, np.integer)):
            return np.array([complex(psi.get(i, 0.0)) for i in range(2**n)], dtype=complex)
        # 索引为二进制字符串（如 "010"）
        if isinstance(k0, str) and set(k0) <= {"0", "1"} and len(k0) == n:
            arr = np.zeros(2**n, dtype=complex)
            for s, amp in psi.items():
                arr[int(s, 2)] = complex(amp)
            return arr
        # 索引为数字字符串（"0", "1", "2", ...）
        if isinstance(k0, str) and k0.isdigit():
            return np.array([complex(psi.get(str(i), 0.0)) for i in range(2**n)], dtype=complex)
    return np.asarray(psi, dtype=complex)


# ============================================================
# 从总态矢量获取相位寄存器边缘分布
# ============================================================
def marginalize_phase_probs_from_statevector(psi: Any,
                                             total_n: int,
                                             phase_qubits: Sequence[int]) -> np.ndarray:
    """
    给定总态矢 psi（包含相位+工作寄存器），
    提取相位寄存器部分的边缘概率分布。

    参数：
        psi          : 态矢量或字典形式
        total_n      : 总量子比特数（相位+工作）
        phase_qubits : 相位寄存器的比特索引（按 LSB→MSB）

    返回：
        phase_probs: 长度为 2^m 的实数组，
                     表示每个相位态的边缘概率。
    """
    amp = _state_to_array(psi, total_n)
    m = len(phase_qubits)
    out = np.zeros(2**m, dtype=float)

    for idx in range(2**total_n):
        # 将全局索引 idx 按小端位序解析
        small = 0
        for i, q in enumerate(phase_qubits):
            bit_q = (idx >> q) & 1
            small |= (bit_q << i)
        out[small] += float(abs(amp[idx])**2)

    return out


# ============================================================
# 从相位寄存器概率分布解码相位 φ ∈ [0,1)
# ============================================================
def decode_phase_from_phase_probs(phase_probs: np.ndarray) -> tuple[str, float]:
    """
    从相位寄存器的概率分布中提取最可能相位。

    参数：
        phase_probs : 长度 2^m 的概率数组，
                      按 LSB→MSB 小端位序排列。

    返回：
        (bits_msb, phi_hat)
        - bits_msb: MSB→LSB 顺序的二进制比特串；
        - phi_hat : 解码得到的相位 ∈ [0,1)。

    过程：
        1. 找到最大概率的索引；
        2. 转为小端二进制字符串 bits_le；
        3. 翻转为 MSB→LSB；
        4. 计算 φ = 0.b0 b1 ... b_{m-1}（二进制小数）。

    示例：
        phase_probs = [0.1, 0.8, 0.05, 0.05]
        → 最大索引=1 → bits_le="01" → bits_msb="10"
        → φ = 0.5 (二进制 0.1)
    """
    t = int(np.log2(len(phase_probs)))
    est_index = int(np.argmax(phase_probs))

    bits_le  = format(est_index, f"0{t}b")   # 小端索引
    bits_msb = bits_le[::-1]                 # 翻转为大端展示

    # 二进制小数 0.b0b1...b_{t-1}
    phi_hat = 0.0
    for i, b in enumerate(bits_msb):
        if b == '1':
            phi_hat += 1.0 / (1 << (i + 1))

    return bits_msb, float(phi_hat)
